import pandas as pd
from sqlalchemy.orm import sessionmaker
from src.db_models.base import engine, Base
from src.db_models.models import Term

# 创建会话
Session = sessionmaker(bind=engine)
session = Session()


def load_excel_to_db(excel_file):
    # 读取 Excel 文件到 DataFrame
    df = pd.read_excel(excel_file)

    # 检查列是否存在
    required_columns = ['CN', 'EN']
    for column in required_columns:
        if column not in df.columns:
            raise ValueError(f"列 {column} 在 Excel 中不存在")

    # 遍历 DataFrame 每一行并插入数据库
    for index, row in df.iterrows():
        term = Term(
            CN=str(row['CN']),  # 转为字符串，确保不会有类型冲突
            EN=str(row['EN']),
        )
        session.add(term)

    # 提交会话到数据库
    session.commit()
    print(f"成功将 {len(df)} 条记录插入数据库！")


if __name__ == "__main__":
    # 数据库初始化，如果表未创建则创建
    Base.metadata.create_all(engine)

    # 调用函数，将 Excel 数据导入数据库
    load_excel_to_db('D:/gitte/project/治国理政第一卷术语库2101.xlsx')